import pandas as pd
import numpy as np
import sys
from sklearn.model_selection import train_test_split  
from sklearn.ensemble import RandomForestRegressor
from tqdm import tqdm
from training.ml_frame.step_rolling.model import *



class Ml_rolling_trin():
    def __init__(self, args) -> None:
        self.args = args

    def consentrate_train(self, X_train, X_out, Y_train, Y_out):
        x_train, _ = data_reshape(X_train)
        x_out, _ = data_reshape(X_out)
        y_train, y_train_shape = data_reshape(Y_train)
        y_out, y_out_shape = data_reshape(Y_out) 
        new_x_train, new_y_train = clean_data(x_train, y_train)
        
        x_out[np.isinf(x_out)] = 0
        x_train[np.isinf(x_train)] = 0
        x_out[np.isnan(x_out)] = 0
        x_train[np.isnan(x_train)] = 0 
        print(f'\t训练开始，特征size:{new_x_train.shape}')     
        if self.args.model_name == 'random_forest':
            model = random_forest_train(new_x_train, new_y_train, self.args)
        if self.args.model_name == 'lightgbm':
            model = lgb_train(new_x_train, new_y_train, self.args)
        if self.args.model_name == 'liner':
            model = liner_train(new_x_train, new_y_train, self.args)
        if self.args.model_name == 'lasso':
            model = lasso_train(new_x_train, new_y_train, self.args)       
        out_pred = model.predict(x_out)
        out_pred[np.isnan(y_out.squeeze(1))] = np.nan
        train_pred = model.predict(x_train) 
        train_pred[np.isnan(y_train.squeeze(1))] = np.nan 

        train_pred = data_reshape_inverse(train_pred, y_train_shape)
        out_pred = data_reshape_inverse(out_pred, y_out_shape)   
        return train_pred, out_pred
    
    def split_train(self,X_train, X_out, Y_train, Y_out):
        train_shape = X_train.shape
        train_pred_lst = []
        out_pred_lst = []
        for code in range(train_shape[1]):
            x_train = X_train[:, code, :]
            y_train = Y_train[:, code, :]
            new_x_train, new_y_train = clean_data(x_train, y_train)
            x_out = X_out[:, code, :]
            y_out = Y_out[:, code, :]           
            x_out[np.isinf(x_out)] = 0
            x_train[np.isinf(x_train)] = 0
            x_out[np.isnan(x_out)] = 0
            x_train[np.isnan(x_train)] = 0 
            print(f'\t训练开始，特征size:{new_x_train.shape}') 
            if self.args.model_name == 'random_forest':
                model = random_forest_train(new_x_train, new_y_train, self.args)
            if self.args.model_name == 'lightgbm':
                model = lgb_train(new_x_train, new_y_train, self.args)
            if self.args.model_name == 'liner':
                model = liner_train(new_x_train, new_y_train, self.args)
            if self.args.model_name == 'lasso':
                model = lasso_train(new_x_train, new_y_train, self.args)
            out_pred = model.predict(x_out)
            out_pred[np.isnan(y_out.squeeze(1))] = np.nan
            train_pred = model.predict(x_train) 
            train_pred[np.isnan(y_train.squeeze(1))] = np.nan
            train_pred_lst.append(train_pred)
            out_pred_lst.append(out_pred)
        train_pred_ary = np.stack(train_pred_lst, axis=-1)
        out_pred_ary = np.stack(out_pred_lst, axis=-1)
        return train_pred_ary, out_pred_ary
    
    def forward_train(self, X_train, X_out, Y_train, Y_out):
        if self.args.model_num == 'one':
            train_pred, out_pred = self.consentrate_train(X_train, X_out, Y_train, Y_out)
        if self.args.model_num == 'codes_num':
            train_pred, out_pred = self.split_train(X_train, X_out, Y_train, Y_out)
        return train_pred, out_pred
            

def clean_data(data_x, data_y):
    data_x[np.isinf(data_x)] = np.nan
    data_y[np.isinf(data_y)] = np.nan
    data = np.concatenate([data_x, data_y], axis=1)
    new_data = data[~np.isnan(data).any(axis=1), :]
    x = new_data[:,:-1]
    y = new_data[:,-1]
    return x, y
    
def data_reshape(array):
    shape = array.shape   
    re_array = array.reshape((-1, shape[-1]), order='F')
    return re_array, shape

def data_reshape_inverse(array, org_shape):
    re_array = array.reshape(org_shape, order='F')
    return re_array.squeeze(2)  
     
def train_eval_resplit(x_train, y_train, x_eval, y_eval):
    x = np.concatenate([x_train, x_eval], axis=0)
    y = np.concatenate([y_train, y_eval], axis=0)
    train_x, eval_x, train_y, eval_y = train_test_split( x, y, test_size=0.3, random_state=2022)
    return train_x, eval_x, train_y, eval_y